package com.zhaopin.trie;

import lombok.Data;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author yin.peide
 * @since 2023-08-18
 */
public class Test {

    /**
     * 插入单词
     * @param rootNode
     * @param word
     */
    public static void insertWord(TrieNode rootNode, String word) {
        TrieNode currentNode = rootNode;
        for (char c : word.toCharArray()) {
            String charStr = String.valueOf(c);
            if (!currentNode.getTireMap().containsKey(charStr)) {
                currentNode.getTireMap().put(charStr, new TrieNode());
            }
            currentNode = currentNode.getTireMap().get(charStr);
        }
        // 在单词的末尾插入 "*" 标记节点
        currentNode.getTireMap().put("*", new TrieNode());
    }

    /**
     * 查询某个单词是否存在
     * @param rootNode
     * @param word
     * @return
     */
    public static boolean searchWord(TrieNode rootNode, String word) {
        TrieNode currentNode = rootNode;
        for (char c : word.toCharArray()) {
            String charStr = String.valueOf(c);
            if (!currentNode.getTireMap().containsKey(charStr)) {
                return false;
            }
            currentNode = currentNode.getTireMap().get(charStr);
        }
        // 检查是否存在 "*" 标记节点，表示找到了完整的单词
        return currentNode.getTireMap().containsKey("*");
    }

    /**
     * 收集所有单词
     * @param currentNode
     * @param prefix
     * @return
     */
    public static List<String> collectWords(TrieNode currentNode, String prefix) {
        List<String> words = new ArrayList<>();
        if (currentNode == null) {
            return words;
        }

        for (String key : currentNode.getTireMap().keySet()) {
            if (key.equals("*")) {
                words.add(prefix);
            } else {
                TrieNode nextNode = currentNode.getTireMap().get(key);
                words.addAll(collectWords(nextNode, prefix + key));
            }
        }
        return words;
    }

    public static void main(String[] args) {
        TrieRoot root = new TrieRoot();

        String[] wordsToIndex = {
                "ace", "act", "bad", "bake", "bat",
                "batter", "cab", "cat", "catnap", "catnip"
        };

        for (String word : wordsToIndex) {
            insertWord(root.getTrieNode(), word);
        }

        // 示例查询操作
        String queryWord = "cat";
        if (searchWord(root.getTrieNode(), queryWord)) {
            System.out.println(queryWord + " found");
        } else {
            System.out.println(queryWord + " not found");
        }

        /*
            收集所有单词
         */
        List<String> allWords = collectWords(root.getTrieNode(), "");
        for (String word : allWords) {
            System.out.print(word+" ");
        }
    }


}

@Data
class TrieRoot {
    private TrieNode trieNode = new TrieNode();
}

@Data
class TrieNode {
    /**
     * key  value 则是节点本身
     */
    private Map<String, TrieNode> tireMap = new HashMap<>();

}